﻿/**

 * Copyright (c) 2015-2017, FastDev 刘强 (fastdev@163.com).

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */


using Mono.Cecil;
using Mono.Cecil.Cil;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using OF.ILModifier.Entity;
using OF.ILModifier.Common;
using OF.ILModifier.Handler;
using OF.ILModifier.Extension;

namespace OF.ILModifier
{
    public class ILBuilder
    {
        private string dllName = null;
        private AssemblyDefinition asm = null;

        private Func<TypeDefinition, bool> targetTypeFilter = null;
        private Func<TypeDefinition, MethodDefinition, Instruction, bool> targetInstructionFilter = null;
        private Func<TypeDefinition, MethodDefinition, bool> targetMethodFilter = null;

        private Func<TypeDefinition, MethodDefinition, Instruction, Instruction[]> beforeInsFunc = null;
        private Func<TypeDefinition, MethodDefinition, Instruction, Instruction[]> afterInsFunc = null;
        private Func<TypeDefinition, MethodDefinition, Instruction, bool> removeInsFunc = null;
        private Tuple<MethodInfo, Func<MethodDefinition, MethodInfo>, MethodInfo> interceptTuple = null;

        private AssemblyHelper asmHelper = null;


        public ILBuilder(string dllName)
        {
            this.dllName = dllName;
            this.asm = AssemblyDefinition.ReadAssembly(dllName);
            this.asmHelper = new AssemblyHelper(this.asm);
        }

        public AssemblyHelper GetAsm()
        {
            return this.asmHelper;
        }

        private void Reset()
        {
            targetTypeFilter = null;
            targetInstructionFilter = null;
            targetMethodFilter = null;
            beforeInsFunc = null;
            afterInsFunc = null;
            removeInsFunc = null;
            interceptTuple = null;
        }

        public ILBuilder SetTargetTypeFilter(Func<TypeDefinition, bool> targetTypeFilter)
        {
            this.targetTypeFilter = targetTypeFilter;
            return this;
        }

        public ILBuilder SetInterceptMethods(MethodInfo beforeMethod, Func<MethodDefinition, MethodInfo> replaceMethod, MethodInfo afterMethod)
        {
            this.interceptTuple = new Tuple<MethodInfo, Func<MethodDefinition, MethodInfo>, MethodInfo>(beforeMethod, replaceMethod, afterMethod);
            return this;
        }
  
        public ILBuilder SetTargetInstructionFilter(Func<TypeDefinition, MethodDefinition, Instruction, bool> targetInstructionFilter)
        {
            this.targetInstructionFilter = targetInstructionFilter;
            return this;
        }

        public ILBuilder SetTargetMethodFilter(Func<TypeDefinition, MethodDefinition, bool> targetMethodFilter)
        {
            this.targetMethodFilter = targetMethodFilter;
            return this;
        }
       
        public ILBuilder SetBeforeInsFunc(Func<TypeDefinition, MethodDefinition, Instruction, Instruction[]> beforeInsFunc)
        {
            this.beforeInsFunc = beforeInsFunc;
            return this;
        }
        
        public ILBuilder SetAfterInsFunc(Func<TypeDefinition, MethodDefinition, Instruction, Instruction[]> afterInsFunc)
        {
            this.afterInsFunc = afterInsFunc;
            return this;
        }


        public ILBuilder SetRemoveInsFunc(Func<TypeDefinition, MethodDefinition, Instruction, bool> removeInsFunc)
        {
            this.removeInsFunc = removeInsFunc;
            return this;
        }


        public ILBuilder Modify(string file)
        {
            RunCore();
            asm.Write(file);
            return this;
        }

        public ILBuilder Visit()
        {
            RunCore();
            return this;
        }

        public ILBuilder Save(string file)
        {
            asm.Write(file);
            return this;
        }

        private void RunCore()
        {
            IEnumerable<TypeDefinition> filterTypes = null;
            if (targetTypeFilter != null)
            {
                filterTypes = asm.MainModule.Types.Where(targetTypeFilter);
            }
            else
            {
                filterTypes = asm.MainModule.Types;
            }
            
            foreach (TypeDefinition type in filterTypes.ToList())
            {
                RunTypeCore(type);
            }
            Reset();
        }

        private MethodInfo GetInterceptReplaceMethod(MethodDefinition method)
        {
            return (interceptTuple.Item2 == null ? null : interceptTuple.Item2(method));
        }
        
        private void RunMethodCore(TypeDefinition type, MethodDefinition method)
        {
            if (method.Body == null || method.Body.Instructions == null)
            {
                return;
            }

            if (interceptTuple != null)
            {
                Console.WriteLine("*******InterceptMethod:" + method.GetFullName());
                new SpecMethodReplaceHandler(this.asmHelper).WrapMethod(type, method, interceptTuple.Item1, GetInterceptReplaceMethod(method), interceptTuple.Item3);
            }

            IEnumerable<Instruction> filterInstructions = null;
            if (targetInstructionFilter != null)
            {
                filterInstructions = Util.Filter(method.Body.Instructions, type, method, targetInstructionFilter);
            }
            else
            {
                filterInstructions = method.Body.Instructions;
            }
            ILProcessor worker = method.Body.GetILProcessor();
            foreach (var ins0 in filterInstructions.ToList())
            {
                if (beforeInsFunc != null)
                {
                    var beforeIns = beforeInsFunc(type, method, ins0);
                    if (beforeIns != null)
                    {
                        foreach (var before in beforeIns)
                        {
                            worker.InsertBefore(ins0, before);
                        }
                    }
                }

                if (afterInsFunc != null)
                {
                    var afterIns = afterInsFunc(type, method, ins0);
                    if (afterIns != null)
                    {
                        foreach (var after in afterIns)
                        {
                            worker.InsertAfter(ins0, after);
                        }
                    }
                }

                if (this.removeInsFunc != null)
                {
                    var isRemove = removeInsFunc(type, method, ins0);
                    if (isRemove)
                    {
                        worker.Remove(ins0);
                    }
                }
            }
        }

        private void RunTypeCore(TypeDefinition type)
        {
            IEnumerable<MethodDefinition> filterMethods = null;
            if (targetMethodFilter != null)
            {
                filterMethods = Util.Filter(type.Methods, type, targetMethodFilter);
            }
            else
            {
                filterMethods = type.Methods;
            }

            foreach (MethodDefinition method in filterMethods.ToList())
            {
                RunMethodCore(type, method);
            }
        }
    }
}